# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: taekwon (originally from BitmexPrivateSubscriber, HuobiFuturesWsPrivateClient)

import enum
import functools
import hashlib
import hmac
import json
import logging
import time
import sys

from tornado.ioloop import IOLoop
from tornado import gen
from tornado.concurrent import Future
from tornado.websocket import websocket_connect, WebSocketClosedError

from coin.base.timestamp import get_timestamp
from coin.exchange.base.rest.task_scheduler import Scheduler
from coin.exchange.base.kr_rest.auth_base import AuthKey

DERIBIT_FUTURES_WS_API_URL = 'wss://www.deribit.com/ws/api/v2'


class OperatorType(enum.Enum):
  SET_HEARTBEAT = 'public/set_heartbeat'
  PING = 'public/test'
  AUTH = 'public/auth'
  SUB = 'private/subscribe'
  UNSUB = 'private/unsubscribe'


class DeribitFuturesWsPrivateClient(object):
  def __init__(self,
               key_file,
               *,
               ioloop=None,
               notification_callback=None,
               close_callback=None,
               on_ready=None,
               timeout_sec=30,
               api_url=DERIBIT_FUTURES_WS_API_URL):
    self._logger = logging.getLogger(__name__)
    self._key = AuthKey.from_file(key_file)

    self._ioloop = ioloop or IOLoop.current()
    self._ws = None
    self._ready = False
    self._stopped = False
    self._timeout = timeout_sec * 1e9
    self._num_timeouts = 0
    self._ws_url = api_url

    self._notification_callback = notification_callback or (lambda x: None)
    self._close_callback = close_callback or (lambda: None)
    self._on_ready = on_ready or (lambda: None)
    self._scheduler = Scheduler()

    self._last_received_message_timestamp = None
    self._pending_futures = {}
    self._subscribed_topics = set()
    self._request_id = int(time.time() * 1000)

    self._init_params = (ioloop, notification_callback, close_callback, timeout_sec)

    self.callbacks = []  # action, table_name, key, data
    self.on_ready_callback = lambda _: None

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  def _get_next_request_id(self):
    self._request_id += 1
    return str(self._request_id)

  @gen.coroutine
  def _check_future_timeout(self):
    timestamp = get_timestamp()
    for req_id, fut in self._pending_futures.copy().items():
      if timestamp - fut.timestamp > self._timeout:
        self._num_timeouts += 1
        fut = self._pending_futures.pop(req_id)
        fut.set_exception(TimeoutError('Websocket request timeout! (%s, %s)' %
                                       (req_id, fut.method)))

  @gen.coroutine
  def _check_connection(self):
    if self._last_received_message_timestamp is None:
      return

    timestamp = get_timestamp()
    diff_time = timestamp - self._last_received_message_timestamp
    if diff_time > self._timeout:
      self._logger.warning('No message received for %s seconds', diff_time / 1e9)
      self.stop()

  @gen.coroutine
  def _ping(self):
    if not self._stopped:
      try:
        self._logger.info('ping')
        response = yield self.send_request(operation=OperatorType.PING)
      except Exception:
        self._logger.exception('Ping error')

  def _gen_auth_msg(self):
    params = {}
    timestamp = str(int(round(time.time() * 1000)))
    nonce = ''.join(reversed(timestamp))
    message = "\n".join([timestamp, nonce, ""])
    signature = hmac.new(bytes(self._key.api_secret, 'utf8'),
                         bytes(message, 'utf8'),
                         digestmod=hashlib.sha256).hexdigest()
    params['params'] = {
        'grant_type': 'client_signature',
        'client_id': self._key.api_key,
        'timestamp': timestamp,
        'nonce': nonce,
        'data': '',
        'signature': signature
    }
    return params

  @gen.coroutine
  def prepare_subscription(self, symbols):
    response = yield self.subscribe_orders(symbols)
    if 'error' in response:
      self.stop()
      return
    else:
      self._logger.info('Subscribe to deribit order %s', symbols)

  @gen.coroutine
  def prepare_subscription_topics(self, topics):
    response = yield self.subscribe_topic(topics)
    if 'error' in response:
      self.stop()
      return
    else:
      self._logger.info('Subscribe to deribit order %s', topics)

  @gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the websocket.', type(self).__name__)
      self._ws = yield websocket_connect(self._ws_url)
      self._ioloop.add_callback(self._run)

      self._logger.info('Send auth info!')
      response = yield self.send_request(operation=OperatorType.AUTH)
      if 'error' in response:
        self._logger.error('Authentication failed! %s', response)
        self.stop()
        return
      else:
        self._logger.info('deribit Futures WS API Successfully authenticated!')

      self._logger.info('Set heartbeat with 40 seconds')
      response = yield self.send_request(operation=OperatorType.SET_HEARTBEAT)
      if 'error' in response:
        self._logger.error('Heartbeat set fail! %s', response)
        self.stop()
        return
      else:
        self._logger.info('Setting heartbeat have done successfully!')

      self._ready = True
      self._on_ready()
      self._logger.info('deribit Futures private client websocket connected.')
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @gen.coroutine
  def _run(self):
    while not self._stopped:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def on_message(self, raw_msg: bytes):
    timestamp = get_timestamp()
    self._last_received_message_timestamp = timestamp
    if self._stopped:
      return

    try:
      data = json.loads(raw_msg)
    except ValueError:
      self._logger.exception('Invalid json string!, %s', raw_msg)
      return

    self._logger.info('Receive ws msg: %s\n', data)

    method = data.get('method', None)
    if method == 'subscription':
      self._notification_callback(data)
    elif method is None:  # something's repsonse. not subscription.
      req_id = data.get('id', None)
      if req_id is None:
        self._logger.error('No request id. %s', data)
        return

      operation = self._pending_futures[req_id].operation
      if operation == OperatorType.SUB:
        topics = data.get('result', None)
        for topic in topics:
          self._subscribed_topics.add(topic)

      fut = self._pending_futures.pop(req_id, None)
      if fut is None:
        self._logger.error('Unexpected message! %s', data)
      else:
        fut.set_result(data)

  def _write_message(self, msg):
    try:
      self._ws.write_message(json.dumps(msg))
    except WebSocketClosedError as e:
      self._logger.error('Write message error, ws is close!')

  def send_request(self, operation, *, topic=None, is_notification=False):
    if self._ws is None:
      self._logger.error('websocket not connected yet! operation=%s, params=%s',
                         operation.value,
                         topic)
      return

    fut = None
    request_id = None
    if not is_notification:
      fut = Future()
      fut.timestamp = get_timestamp()
      fut.operation = operation
      request_id = self._get_next_request_id()
      self._pending_futures[request_id] = fut

    request_msg = {'jsonrpc': '2.0', 'id': request_id}
    if operation == OperatorType.AUTH:
      request_msg['method'] = operation.value
      request_msg.update(self._gen_auth_msg())
    elif operation == OperatorType.SET_HEARTBEAT:
      request_msg['method'] = operation.value
      request_msg['params'] = {'interval': 40}
    elif operation == OperatorType.PING:
      request_msg['method'] = operation.value
    else:
      assert topic is not None
      request_msg['method'] = operation.value
      request_msg['params'] = {'channels': topic}

    self._logger.info('send request: %s\n', request_msg)
    self._ioloop.add_callback(self._write_message, request_msg)
    return fut

  def subscribe_topic(self, topic):
    operation = OperatorType.SUB
    fut = self.send_request(operation, topic=topic)
    return fut

  def unsubscribe_topic(self, topic):
    operation = OperatorType.UNSUB
    fut = self.send_request(operation, topic=topic)
    return fut

  def subscribe_orders(self, symbols):
    topic = ['user.orders.{symbol}.raw'.format(symbol=symbol) for symbol in symbols]
    fut = self.subscribe_topic(topic)
    return fut

  def unsubscribe_orders(self, symbols):
    topic = ['user.orders.{symbol}.raw'.format(symbol=symbol) for symbol in symbols]
    fut = self.unsubscribe_topic(topic)
    return fut

  def start(self):
    self._ioloop.add_callback(self._connect)
    self._scheduler.add_periodic_task(20, self._ping)
    self._scheduler.add_periodic_task(1, self._check_future_timeout)
    self._scheduler.add_periodic_task(10, self._check_connection)
    self._scheduler.start()

  def stop(self):
    if self._stopped:
      return
    self._scheduler.stop()
    self._ready = False
    self._ioloop.add_callback(self._async_stop)
    self._logger.info('deribit ws private client (%s) stopped!', id(self))

  @gen.coroutine
  def _async_stop(self):
    for topic in self._subscribed_topics:
      response = yield self.unsubscribe_topic(topic)
      self._logger.info('Topic (%s) unsubscribed!\n%s', topic, response)

    self._stopped = True
    if self._ws is not None:
      self._ws.close()
      self._ws = None
    yield gen.sleep(10)
    self._close_callback()

  def clone(self):
    return type(self)(*self._init_params)


class TestWsPrivateClient(object):
  def __init__(self, key_filepath, ioloop=None):
    self._key_filepath = key_filepath
    self._ioloop = ioloop
    self._ws_private_client = None

  def _prepare_ws_private_client(self, start=False):
    if self._ws_private_client is not None:
      self._ws_private_client.stop()

    self._ws_private_client = DeribitFuturesWsPrivateClient(
        self._key_filepath,
        ioloop=self._ioloop,
        close_callback=functools.partial(self._prepare_ws_private_client, True),
        on_ready=self._on_ws_private_client_ready,
    )

    if start:
      self._ws_private_client.start()

  def _on_ws_private_client_ready(self):
    topics = ['user.portfolio.btc', 'user.orders.any.BTC.raw', 'deribit_price_index.btc_usd']
    self._ws_private_client.prepare_subscription_topics(topics)

  def start(self):
    self._prepare_ws_private_client(start=True)


def main(argv):
  key_filepath = '/home/taekwon/deribit_api_key.json'

  ioloop = IOLoop.current()
  test_client = TestWsPrivateClient(key_filepath, ioloop)
  test_client.start()
  ioloop.start()


if __name__ == '__main__':
  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  sys.exit(main(sys.argv))
